import numpy as np
import tensorflow as tf
import t3f
import cv2
import math


IMG_WIDTH = 512
IMG_HEIGHT = 384
l_width_modes = [4,4,4,8]
l_height_modes = [6,4,4,4]
l_tt_ranks = [24,24,24]


def tt_Hadamard(tt_left, tt_right):
	# 哈达玛积，t3f.multiply是有问题的，自己定义一个
	cores_left = tt_left._tt_cores
	cores_right = tt_right._tt_cores

	n = len(cores_left)
	new_cores = []
	for i in range(n):
		new_cores.append(cores_left[i] * cores_right[i])

	return t3f.TensorTrain(new_cores)


def read_img(name):
	# 读取为单通道灰度图像，RGB是三通道图像，如果想complete彩图，就按照R、G、B三个通道分别做
	img_data = cv2.imread(name, flags = cv2.IMREAD_GRAYSCALE)
	if img_data.shape[0] != IMG_HEIGHT or img_data.shape[1] != IMG_WIDTH:
		return None
	return img_data


def riemannian(ph_groud_truth):
	# 固定随机数，保证random_normal得到的noise和random_uniform得到的sparsity_mask是固定的，不会在每次run时变化
	tf.set_random_seed(0)

	# ph_groud_truth代表原图，noisy_ground_truth代表可观测的加了噪声的实际原图
	noise = 1e-2 * tf.get_variable('noise', initializer = tf.random_normal(ph_groud_truth.shape), trainable = False)
	noisy_ground_truth = ph_groud_truth + noise

	# mask，表示遮蔽原图的缺失部分，这里意味着75%的像素可见，其余被遮蔽
	sparsity_mask = tf.cast(tf.random_uniform(ph_groud_truth.shape) <= 0.75, tf.float32)
	sparsity_mask = tf.get_variable('sparsity_mask', initializer = sparsity_mask, trainable=False)
	tt_sparsity_mask = t3f.to_tt_matrix(sparsity_mask, [l_height_modes, l_width_modes], np.min(np.array([IMG_WIDTH, IMG_HEIGHT])), 1e-4)

	# 可观测到的所有像素数量
	observed_total = tf.reduce_sum(sparsity_mask)

	# 原图所有像素数量
	total = np.prod(np.array(ph_groud_truth.shape.as_list()))

	# 实际的观测值
	sparse_observation = noisy_ground_truth * sparsity_mask

	# 实际观测值作为输入转化为TT张量
	tt_sparse_observation = t3f.to_tt_matrix(sparse_observation, [l_height_modes, l_width_modes], np.min(np.array([IMG_WIDTH, IMG_HEIGHT])), 1e-4)

	# 优化目标张量(流形上的点)
	initializer_manifold = t3f.glorot_initializer([l_height_modes, l_width_modes], tt_rank = [1] + l_tt_ranks + [1])
	tt_manifold = t3f.get_variable('manifold', initializer = initializer_manifold)

	# 训练loss，SGD
	grad = tt_manifold * tt_sparsity_mask - tt_sparse_observation
	riemannian_grad = t3f.riemannian.project(grad, tt_manifold)
	train_step = t3f.assign(tt_manifold, t3f.round(tt_manifold - 0.001 * riemannian_grad, max_tt_rank = [1] + l_tt_ranks + [1]))
	loss = 1.0 / observed_total * t3f.frobenius_norm_squared(tt_manifold * tt_sparsity_mask - tt_sparse_observation)

	# 测试loss
	test_loss = 1.0 / total * tf.reduce_sum((t3f.full(tt_manifold) - ph_groud_truth)**2)

	return [loss, test_loss], sparse_observation, t3f.full(tt_manifold), train_step.op


def main(_):
	# 读图，得到其数据
	img_data = read_img('Tulips.jpg')
	if img_data is None:
		print('Image sizes are not matched.')
		return

	# 构建loss和test_loss，前者是在有mask的前提下，计算mask后的估计图和mask后的噪声图之间的误差，后者是测试无mask无噪声的估计图和原图之间的误差
	ph_groud_truth = tf.placeholder(dtype = tf.float32, shape = img_data.shape, name = 'ph_groud_truth')
	l_losses, ob, es, train_step = riemannian(ph_groud_truth)

	# 定义session，初始化全部的variables(tf.get_variable和t3f.get_variable定义的变量)
	sess = tf.Session(config = tf.ConfigProto(allow_soft_placement = True, device_count = {'GPU': 0}))
	sess.run(tf.global_variables_initializer())

	# 训练过程，训练n次
	n = 10000
	for i in range(n):
		# 注意，run时，ph_groud_truth才有了来自img_data的数据
		_, tr_loss_v, test_loss_v = sess.run([train_step] + l_losses, feed_dict = {ph_groud_truth : img_data})
		if math.isnan(tr_loss_v) or math.isnan(test_loss_v):
			print('NaN! The norm of estimated is %f.\n' % sess.run(tf.norm(es), feed_dict = {ph_groud_truth : img_data}))
			break
		if i % 100 == 0:
			print('Step %d/%d. Current train loss is %f, test loss is %f.\n' % (i, n, tr_loss_v, test_loss_v))

	# 保存原图
	cv2.imwrite('ori_mf.jpg', img_data)

	# 保存可观测的图
	ob_data = sess.run(ob, feed_dict = {ph_groud_truth : img_data})
	cv2.imwrite('observed_mf.jpg', ob_data)

	# 保存估计图
	es_data = sess.run(es)
	cv2.imwrite('estimated_mf.jpg', es_data)


if __name__ == '__main__':
    tf.app.run()
